package com.boulderai.mcp.example;

import com.boulderai.mcp.client.McpClient;
import com.boulderai.mcp.client.McpClientFactory;
import com.boulderai.mcp.model.*;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.node.ObjectNode;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.List;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.TimeUnit;

/**
 * 演示如何连接到本地MCP服务器的示例
 * 
 * 使用方法：
 * 1. 启动你的本地MCP服务器（例如在端口8080）
 * 2. 运行这个示例程序
 * 3. 客户端会自动连接并注册到服务器
 */
public class LocalServerExample {
    private static final Logger logger = LoggerFactory.getLogger(LocalServerExample.class);
    private static final ObjectMapper objectMapper = new ObjectMapper();
    
    public static void main(String[] args) {
        // 默认连接到本地服务器
        String serverUri = "ws://localhost:8080";
        
        // 如果提供了命令行参数，使用指定的服务器地址
        if (args.length > 0) {
            serverUri = args[0];
        }
        
        logger.info("正在连接到本地MCP服务器: {}", serverUri);
        
        try {
            // 创建并连接到MCP服务器
            McpClient client = connectToLocalServer(serverUri);
            
            // 演示各种操作
            demonstrateServerOperations(client);
            
            // 保持连接一段时间
            logger.info("保持连接30秒，然后断开...");
            Thread.sleep(30000);
            
            // 断开连接
            client.disconnect();
            logger.info("已断开与服务器的连接");
            
        } catch (Exception e) {
            logger.error("连接本地服务器时发生错误: {}", e.getMessage(), e);
        }
    }
    
    /**
     * 连接到本地MCP服务器
     */
    private static McpClient connectToLocalServer(String serverUri) throws Exception {
        // 创建客户端
        McpClient client = McpClientFactory.createWebSocketClient();
        
        // 连接到服务器
        CompletableFuture<Void> connectFuture = client.connect(serverUri);
        connectFuture.get(10, TimeUnit.SECONDS); // 等待连接完成
        
        logger.info("成功连接到MCP服务器: {}", serverUri);
        
        // 初始化会话
        ObjectNode clientInfo = objectMapper.createObjectNode();
        clientInfo.put("name", "LocalServerExample");
        clientInfo.put("version", "1.0.0");
        ObjectNode capabilities = objectMapper.createObjectNode();
        capabilities.put("tools", true);
        capabilities.put("resources", true);
        clientInfo.set("capabilities", capabilities);
        
        CompletableFuture<JsonNode> initFuture = client.initialize(clientInfo);
        JsonNode initResponse = initFuture.get(5, TimeUnit.SECONDS);
        
        logger.info("会话初始化成功");
        if (initResponse.has("serverInfo")) {
            JsonNode serverInfo = initResponse.get("serverInfo");
            logger.info("服务器信息: {} v{}", 
                serverInfo.get("name").asText(),
                serverInfo.get("version").asText());
        }
        if (initResponse.has("capabilities")) {
            logger.info("服务器能力: {}", initResponse.get("capabilities"));
        }
        
        return client;
    }
    
    /**
     * 演示服务器操作
     */
    private static void demonstrateServerOperations(McpClient client) {
        try {
            // 1. 列出可用工具
            logger.info("\n=== 获取可用工具 ===");
            CompletableFuture<List<Tool>> toolsFuture = client.listTools();
            List<Tool> tools = toolsFuture.get(5, TimeUnit.SECONDS);
            
            if (tools.isEmpty()) {
                logger.info("服务器没有提供任何工具");
            } else {
                logger.info("服务器提供的工具:");
                for (Tool tool : tools) {
                    logger.info("  - {}: {}", tool.getName(), tool.getDescription());
                }
                
                // 尝试调用第一个工具
                Tool firstTool = tools.get(0);
                callToolExample(client, firstTool);
            }
            
            // 2. 列出可用资源
            logger.info("\n=== 获取可用资源 ===");
            CompletableFuture<List<Resource>> resourcesFuture = client.listResources();
            List<Resource> resources = resourcesFuture.get(5, TimeUnit.SECONDS);
            
            if (resources.isEmpty()) {
                logger.info("服务器没有提供任何资源");
            } else {
                logger.info("服务器提供的资源:");
                for (Resource resource : resources) {
                    logger.info("  - {}: {}", resource.getUri(), resource.getDescription());
                }
                
                // 尝试读取第一个资源
                Resource firstResource = resources.get(0);
                readResourceExample(client, firstResource);
            }
            
            // 3. 发送自定义请求示例
            logger.info("\n=== 发送自定义请求 ===");
            try {
                ObjectNode params = objectMapper.createObjectNode();
                CompletableFuture<JsonNode> customFuture = client.sendRequest("ping", params);
                JsonNode customResponse = customFuture.get(5, TimeUnit.SECONDS);
                logger.info("自定义请求响应: {}", customResponse);
            } catch (Exception e) {
                logger.info("服务器不支持ping请求: {}", e.getMessage());
            }
            
        } catch (Exception e) {
            logger.error("演示服务器操作时发生错误: {}", e.getMessage(), e);
        }
    }
    
    /**
     * 调用工具示例
     */
    private static void callToolExample(McpClient client, Tool tool) {
        try {
            logger.info("\n=== 调用工具: {} ===", tool.getName());
            
            // 根据工具类型准备参数
            ObjectNode arguments = objectMapper.createObjectNode();
            if ("echo".equals(tool.getName())) {
                arguments.put("message", "Hello from LocalServerExample!");
            } else if ("add".equals(tool.getName())) {
                arguments.put("a", 10);
                arguments.put("b", 20);
            }
            
            CompletableFuture<JsonNode> future = client.callTool(tool.getName(), arguments);
            JsonNode response = future.get(5, TimeUnit.SECONDS);
            
            logger.info("工具调用结果: {}", response);
            
        } catch (Exception e) {
            logger.error("调用工具时发生错误: {}", e.getMessage(), e);
        }
    }
    
    /**
     * 读取资源示例
     */
    private static void readResourceExample(McpClient client, Resource resource) {
        try {
            logger.info("\n=== 读取资源: {} ===", resource.getUri());
            
            CompletableFuture<JsonNode> future = client.readResource(resource.getUri());
            JsonNode response = future.get(5, TimeUnit.SECONDS);
            
            logger.info("资源内容: {}", response);
            
        } catch (Exception e) {
            logger.error("读取资源时发生错误: {}", e.getMessage(), e);
        }
    }
}